#!/usr/bin/env python

# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Test a Fast R-CNN network on an image database."""

import _init_paths
from fast_rcnn.dj_volmat_test_luna_kaggle import test_net
from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list
from datasets.factory import get_imdb
import caffe
import argparse
import pprint
import time, os, sys

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Test a Fast R-CNN network')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU id to use',
                        default=0, type=int)
    parser.add_argument('--def', dest='prototxt',
                        help='prototxt file defining the network',
                        default=None, type=str)
    parser.add_argument('--net', dest='caffemodel',
                        help='model to test',
                        default=None, type=str)
    parser.add_argument('--fold_idx', dest='fold_idx',
                        help='vol idx',
                        default=None, type=int)
    parser.add_argument('--vol_dir', dest='vol_dir',
                        help='dir of test vol.mat',
                        default=None, type=str)
    parser.add_argument('--candi_dir', dest='candi_dir',
                        help='dir of test candi.mat',
                        default=None, type=str)

    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file', default=None, type=str)
    parser.add_argument('--wait', dest='wait',
                        help='wait until net file exists',
                        default=True, type=bool)
    parser.add_argument('--imdb', dest='imdb_name',
                        help='dataset to test',
                        default='voc_2007_test', type=str)
    parser.add_argument('--comp', dest='comp_mode', help='competition mode',
                        action='store_true')
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument('--vis', dest='vis', help='visualize detections',
                        action='store_true')
    parser.add_argument('--num_dets', dest='max_per_image',
                        help='max number of detections per image',
                        default=100, type=int)

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    return args

'''
--def models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt
--cfg experiments/cfgs/faster_rcnn_end2end.yml
--net output/faster_rcnn_end2end/voc_2007_trainval/vgg16_faster_rcnn_iter_60000.caffemodel
--fold_idx 8
--candi_dir output/candis_kaggle_beni_60000
--gpu 0
'''
def test_process(gpu_id,
                 caffemodel,
                 candi_dir,
                 fold_idx,
                 start_scan_idx,
                 end_scan_idx,
                 fold_list_path,
                 kaggle_vol_path,
                 lidc_vol_path,
                 spie_vol_path,
                 stage2_vol_path,
                 prototxt='models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt',
                 flag_wait=True
                 ):
    if stage2_vol_path is None and fold_idx > 8:
        print 'stage2_vol_path should be assigned or fold_idx should < 9'
        return

    cfg_from_file('experiments/cfgs/faster_rcnn_end2end.yml')
    cfg.GPU_ID = gpu_id
    print('Using config:')
    pprint.pprint(cfg)

    while not os.path.exists(caffemodel) and flag_wait:
        print('Waiting for {} to exist...'.format(caffemodel))
        time.sleep(10)
    if not os.path.exists(candi_dir) and (candi_dir is not None):
        os.makedirs(candi_dir)

    caffe.set_mode_gpu()
    caffe.set_device(gpu_id)
    #caffe.set_mode_cpu()

    net = caffe.Net(prototxt, caffemodel, caffe.TEST)
    net.name = os.path.splitext(os.path.basename(caffemodel))[0]

    test_net(net,
             fold_idx=fold_idx,
             start_scan_idx=start_scan_idx,
             end_scan_idx=end_scan_idx,
             fold_list_path=fold_list_path,
             candi_dir=candi_dir,
             kaggle_vol_path=kaggle_vol_path,
             lidc_vol_path=lidc_vol_path,
             spie_vol_path=spie_vol_path,
             stage2_vol_path=stage2_vol_path,
             max_per_image=100,
             thresh=0.05,
             vis=False)


    pass


if __name__ == '__main__':
    args = parse_args()

    print('Called with args:')
    print(args)

    if args.cfg_file is not None:
        cfg_from_file(args.cfg_file)
    if args.set_cfgs is not None:
        cfg_from_list(args.set_cfgs)

    cfg.GPU_ID = args.gpu_id

    print('Using config:')
    pprint.pprint(cfg)

    while not os.path.exists(args.caffemodel) and args.wait:
        print('Waiting for {} to exist...'.format(args.caffemodel))
        time.sleep(10)
    if not os.path.exists(args.candi_dir) and (args.candi_dir is not None):
        os.mkdir(args.candi_dir)

   # caffe.set_mode_gpu()
    #caffe.set_device(args.gpu_id)
    caffe.set_mode_cpu()

    net = caffe.Net(args.prototxt, args.caffemodel, caffe.TEST)
    net.name = os.path.splitext(os.path.basename(args.caffemodel))[0]

    test_net(net, fold_idx=args.fold_idx, candi_dir=args.candi_dir,
             max_per_image=args.max_per_image, vis=args.vis)
